#! /usr/bin/env python
#-*- coding=utf-8 -*-
import rospy
from sensor_msgs.msg import Image, PointCloud2, Imu, NavSatFix
import cv2
import os
from cv_bridge import CvBridge

import numpy as np
from std_msgs.msg import Header
import sensor_msgs.point_cloud2 as point_cloud2

from visualization_msgs.msg import Marker, MarkerArray
from geometry_msgs.msg import Point
import tf

import pandas as pd
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

from kitti_utils import *


DATA_PATH= '/home/lincent/dataset/KITTI/2011_09_26_drive_0005_sync/2011_09_26/2011_09_26_drive_0005_sync/'
FRAME_ID= 'map'

COLOR_DIC= {'Car':(255,0,0), 'Pedestrian':(0,255,255), 'Cyclist':(40,40,255)}  # 类别-颜色对应的字典 

#OXTS 文件
IMU_COLUMN_NAMES= ['lat', 'lon', 'alt', 'roll', 'pitch', 'yaw', 'vn', 've', 'vf', 'vl', 'vu',
                   'ax', 'ay', 'az', 'af', 'al', 'au', 'wx', 'wy', 'wz', 'wf', 'wl', 'wu', 
                   'posacc', 'velacc', 'navstat', 'numsats', 'posmode', 'velmode', 'orimode']
# trcaking label 文件
TRACKING_COLUMN_NAMES= ['frame', 'track_id', 'type', 'truncated', 'occluded', 'alpha', 
                       'bbox_left', 'bbox_top', 'bbox_right', 'bbox_bottom', 'height', 'width', 
                       'length', 'pos_x', 'pos_y', 'pos_z', 'rot_y']

def image_publish(cam_pub, bridge, img ,boxes, types):
    for ty, box in zip(types, boxes):    # zip 操作 合并两个tuple
        top_left= int(box[0]), int(box[1])  # tuple
        bottom_right= int(box[2]), int(box[3])
        cv2.rectangle(img, top_left, bottom_right, COLOR_DIC[ty], 3)
    cam_pub.publish(bridge.cv2_to_imgmsg(img, 'bgr8'))


def point_cloud_publish(pc_pub, ponit_cloud):
    header= Header()
    header.stamp= rospy.Time.now()
    header.frame_id= FRAME_ID
    pc_pub.publish(point_cloud2.create_cloud_xyz32(header, ponit_cloud[:, : 3]))

def view_publish(view_pub):
    marker= Marker()
    marker.header.frame_id= FRAME_ID
    marker.header.stamp= rospy.Time.now()
    marker.id= 0
    marker.action= Marker.ADD
    marker.lifetime= rospy.Duration()
    marker.type= Marker.LINE_STRIP

    marker.color.r= 0.0
    marker.color.g= 1.0
    marker.color.b= 0.0
    marker.color.a= 1.0
    marker.scale.x= 0.2

    marker.points= []
    marker.points.append(Point(10, 10, 0))
    marker.points.append(Point(0,0,0))
    marker.points.append(Point(10,-10,0))

    view_pub.publish(marker)

def car_model_publish(car_model_pub):
    mesh_marker= Marker()
    mesh_marker.header.stamp= rospy.Time.now()
    mesh_marker.header.frame_id= FRAME_ID

    mesh_marker.id= -1
    mesh_marker.lifetime= rospy.Duration()
    mesh_marker.type= Marker.MESH_RESOURCE
    #mesh_marker.mesh_resource= 'package://kitti/India_City_Car/91240_India_City_Car.dae'
    mesh_marker.mesh_resource= 'package://kitti/1553_Peugeot_407_Rally_Car_Concept/1553_Peugeot_407_Rally_Car_Concept.dae'

    mesh_marker.pose.position.x= 0.0
    mesh_marker.pose.position.y= 0.0
    mesh_marker.pose.position.z= 0.0  #-1.73

    q= tf.transformations.quaternion_from_euler(np.pi/2, 0 ,np.pi/2)
    mesh_marker.pose.orientation.x= q[0]
    mesh_marker.pose.orientation.y= q[1]
    mesh_marker.pose.orientation.z= q[2]
    mesh_marker.pose.orientation.w= q[3]

    mesh_marker.color.r= 1.0
    mesh_marker.color.g= 1.0
    mesh_marker.color.b= 1.0
    mesh_marker.color.a= 1.0

    mesh_marker.scale.x= 0.006
    mesh_marker.scale.y= 0.006
    mesh_marker.scale.z= 0.006
    
    car_model_pub.publish(mesh_marker)

def imu_publish(imu_pub, imu_data):
    imu= Imu()
    imu.header.stamp= rospy.Time.now()
    imu.header.frame_id=  FRAME_ID

    q= tf.transformations.quaternion_from_euler(float(imu_data.roll), float(imu_data.pitch), float(imu_data.yaw))
    imu.orientation.x= q[0]
    imu.orientation.y= q[1]
    imu.orientation.z= q[2]
    imu.orientation.w= q[3]
    imu.linear_acceleration.x= imu_data.af
    imu.linear_acceleration.y= imu_data.al
    imu.linear_acceleration.z= imu_data.au
    imu.angular_velocity.x= imu_data.wf
    imu.angular_velocity.z= imu_data.wl
    imu.angular_velocity.z- imu_data.wu

    imu_pub.publish(imu)

def gps_publish(gps_pub, gps_data):
    gps= NavSatFix()
    gps.header.stamp= rospy.Time.now()
    gps.header.frame_id= FRAME_ID
    gps.latitude= gps_data.lat
    gps.longitude= gps_data.lon
    gps.altitude= gps_data.alt

    gps_pub.publish(gps)

# box corner在02 相机坐标系的坐标
def compute_3d_bbox_image02(h, w, l, x, y, z, yaw):
    R =np.array([[np.cos(yaw), 0, np.sin(yaw)], [0,1,0], [-np.sin(yaw), 0, np.cos(yaw)]])
    x_corners= [l/2, l/2, -l/2, -l/2, l/2, l/2, -l/2, -l/2]
    y_corners= [0,0,0,0, -h, -h, -h, -h]
    z_corners= [w/2, -w/2, -w/2, w/2, w/2, -w/2, -w/2, w/2]
    corners_3d_cam2= np.dot(R, np.vstack([x_corners, y_corners, z_corners]))
    corners_3d_cam2+= np.vstack([x,y,z])
    return corners_3d_cam2

# box corner的连接关系
LINES= [[0,1],[1,2],[2,3],[3,0],  #lower
        [4,5],[5,6],[6,7],[7,4],  #upper
        [0,4],[1,5],[2,6],[3,7],  #between lower and upper
        [0,5],[1,4]]             #front x

def bbox_3d_publish(bbox_3d_pub, boxes_corners_3d_velo, types, track_ids):
    marker_array= MarkerArray()
    for i, box_corners_3d_velo in enumerate(boxes_corners_3d_velo):  # 返回第几个box
        marker= Marker()
        marker.header.stamp= rospy.Time.now()
        marker.header.frame_id= FRAME_ID
        marker.id= i
        marker.action= Marker.ADD
        marker.lifetime= rospy.Duration(0.1)   # 根据帧率设置每个box的存在时间 0.1s
        marker.type= Marker.LINE_LIST

        b,g,r= COLOR_DIC[types[i]]
        marker.color.r= r/255.0
        marker.color.g= g/255.0
        marker.color.b= b/255.0
        marker.color.a= 1.0
        marker.scale.x= 0.1

        marker.points= []
        for l in LINES:
            p1= box_corners_3d_velo[l[0]]
            marker.points.append(Point(p1[0], p1[1], p1[2]))
            p2= box_corners_3d_velo[l[1]]
            marker.points.append(Point(p2[0], p2[1], p2[2]))
        marker_array.markers.append(marker)

        ## id text
        text_marker= Marker()
        text_marker.header.stamp= rospy.Time.now()
        text_marker.header.frame_id= FRAME_ID
        text_marker.id= i+ 1000         #避免id重复
        text_marker.action= Marker.ADD
        text_marker.lifetime= rospy.Duration(0.1)   # 根据帧率设置每个box的存在时间 0.1s
        text_marker.type= Marker.TEXT_VIEW_FACING   #text始终朝向观察方向

        text_marker.text= str(track_ids[i])

        p= np.mean(box_corners_3d_velo, axis= 0)  # id的位置
        text_marker.pose.position.x= p[0]
        text_marker.pose.position.y= p[1]
        text_marker.pose.position.z= p[2]+1.5
        text_marker.scale.x= 1
        text_marker.scale.y= 1
        text_marker.scale.z= 1

        text_marker.color.r= r/255.0
        text_marker.color.g= g/255.0
        text_marker.color.b= b/255.0
        text_marker.color.a= 1.0
        marker_array.markers.append(text_marker)
    

    bbox_3d_pub.publish(marker_array)

#from collections import deque
class Object():
    def __init__(self, center):
        self.locations= []
        self.locations.append(center)
        #self.locations= deque(maxlen=20)
    def update(self, center, dis, yaw_change):
        for i in range(len(self.locations)):
            x0, y0= self.locations[i]
            x1= x0* np.cos(yaw_change) + y0* np.sin(yaw_change) -dis
            y1= -x0* np.sin(yaw_change) + y0* np.cos(yaw_change)
            self.locations[i]= np.array([x1,y1])
        if center is not None:
            self.locations.append(center)
        #self.locations= self.locations[:20]  #只取前20个 一段轨迹
        #self.locations.appendleft(np.array([0,0]))
    def reset(self):
        self.locations= []
        #self.locations= deque(maxlen=20)

def location_publish(loc_pub, tracker, centers):
    markerArray= MarkerArray()
    for track_id in centers:  #只画出当前帧的
        marker= Marker()
        marker.header.stamp= rospy.Time.now()
        marker.header.frame_id= FRAME_ID

        marker.id= track_id
        marker.action= Marker.ADD
        marker.lifetime= rospy.Duration(0.1)
        marker.type= Marker.LINE_STRIP
        
        marker.color.r= 1.0
        marker.color.g= 0.2
        marker.color.b= 0.2
        marker.color.a= 1.0
        marker.scale.x= 0.2
        marker.points= []
        for p in tracker[track_id].locations:
            marker.points.append(Point(p[0], p[1], 0))
        markerArray.markers.append(marker)
    loc_pub.publish(markerArray)

if __name__=='__main__':
    rospy.init_node('kitti_node', anonymous= True)
 
    cam_pub= rospy.Publisher('kitti_cam', Image, queue_size= 10)
    pc_pub= rospy.Publisher('kitti_point_cloud', PointCloud2, queue_size= 10)
    view_pub= rospy.Publisher('kitti_cam_view', Marker, queue_size= 10)
    car_model_pub= rospy.Publisher('kitti_car_model', Marker, queue_size= 10)
    imu_pub= rospy.Publisher('kitti_imu', Imu, queue_size= 10)
    gps_pub= rospy.Publisher('kitti_gps', NavSatFix, queue_size= 10)
    bbox_3d_pub= rospy.Publisher('kitti_3d_bbox', MarkerArray, queue_size= 10)
    loc_pub= rospy.Publisher('kitti_location', MarkerArray, queue_size= 20)

    
    bridge= CvBridge()   
    calib= Calibration(os.path.join(DATA_PATH, '2011_09_26_calib/2011_09_26/'), from_video= True)
    tracking_df= pd.read_csv(os.path.join(DATA_PATH, 'data_tracking_label_2/training/label_02/0000.txt'), 
                             header= None, sep=' ')
    tracking_df.columns= TRACKING_COLUMN_NAMES

    tracker= {}
    previous_imu= None
    frame= 0

    rate= rospy.Rate(10)
    while not rospy.is_shutdown():
        #point cloud
        ponit_cloud= np.fromfile(os.path.join(DATA_PATH, 'velodyne_points/data/%010d.bin'% frame), 
                                dtype=np.float32).reshape(-1,4)
        point_cloud_publish(pc_pub, ponit_cloud)
        #image  2d bbox
        tracking_df.loc[tracking_df.type.isin(['Tram', 'Van', 'Truck']), 'type']= 'Car'
        tracking_df= tracking_df[tracking_df.type.isin(['Car', 'Pedestrian', 'Cyclist'])]
        df= tracking_df[tracking_df.frame== frame]
        boxes= np.array(df[['bbox_left', 'bbox_top', 'bbox_right', 'bbox_bottom']])
        types= np.array(df['type'])
        track_ids= np.array(df['track_id'])

        img= cv2.imread(os.path.join(DATA_PATH, 'image_02/data/%010d.png'% frame))
        image_publish(cam_pub, bridge, img, boxes, types)
        #3d bbox
        bbox_3d= np.array(df[['height', 'width', 'length', 'pos_x', 'pos_y', 'pos_z', 'rot_y']])
        boxes_corners_3d_velo= []
        for box in bbox_3d:
            corners_3d_cam2= compute_3d_bbox_image02(*box)
            corners_3d_velo= calib.project_rect_to_velo(corners_3d_cam2.T)
            boxes_corners_3d_velo += [corners_3d_velo]
        bbox_3d_publish(bbox_3d_pub, boxes_corners_3d_velo, types, track_ids)
        #view
        view_publish(view_pub)
        #car model
        car_model_publish(car_model_pub)
        # imu
        df= pd.read_csv(os.path.join(DATA_PATH, 'oxts/data/%010d.txt'% frame), header= None, sep=' ')
        df.columns= IMU_COLUMN_NAMES
        imu_publish(imu_pub, df)
        # gps
        gps_publish(gps_pub, df)
        # loc
        centers= {}
        centers[-1]= np.array([0,0])
        for track_id, box_corners_3d_velo in zip(track_ids, boxes_corners_3d_velo): 
            centers[track_id]= np.mean(box_corners_3d_velo, axis= 0)[:2]

        if previous_imu is None:
            for track_id in centers:
                tracker[track_id]= Object(centers[track_id])
        else:
            dis= 0.1* np.linalg.norm(df[['vf', 'vl']])
            yaw_change= float(df.yaw- previous_imu.yaw)
            for track_id in centers:
                if track_id in tracker:
                    tracker[track_id].update(centers[track_id], dis, yaw_change)
                else:
                    tracker[track_id]= Object(centers[track_id])
            for track_id in tracker:        
                if track_id not in centers:
                    tracker[track_id].update(None, dis, yaw_change)

        previous_imu= df
        location_publish(loc_pub, tracker, centers)

        rospy.loginfo("published")
        rate.sleep()

        frame+= 1
        frame %=154 
        if(frame ==0 ):
            for track_id in tracker:
                tracker[track_id].reset()